#!/usr/bin/env python

from numpy import arange
from operator import itemgetter
from pylab import exp, hold, plot, show
from random import random
import sys

def logistic_func(a):
    return lambda x: a * x * (1 - x)

def iterate(f, start, n):
    ans = [start]
    for i in range(n):
        ans.append(f(ans[-1]))
    return ans

def attractors(f, start):
    orbit = iterate(f, start, 500)
    return orbit[200:-1]

def plot_attractors(param, f, viewfunc, start):
    orbit = attractors(f, start)
    plot([param for x in orbit], [viewfunc(x) for x in orbit], 'r,')

def bi_diagram(f_constructor, param_range, viewfunc, start):
    hold(True)
    for param in param_range:
        plot_attractors(param, f_constructor(param), viewfunc, start)
    show()

def logistic_diagram():
    bi_diagram(logistic_func, arange(1, 4, 0.01), lambda x: x, random())

def beetles(mu_a, state):
    """Lab Visit I from Alligood et al., Chaos
    """
    c_el = 0.012
    c_ea = 0.009
    c_pa = 0.004
    mu_l = 0.267
    mu_p = 0
    b = 7.48
    L, P, A = state
    L_new = b * A * exp((-c_ea * A) - (c_el * L))
    P_new = L * (1 - mu_l)
    A_new = P * (1 - mu_p) * exp(-c_pa * A) + A * (1 - mu_a)
    return (L_new, P_new, A_new)

def beetle_func(mu_a):
    return lambda s: beetles(mu_a, s)

def beetle_diagram():
    print("Running beetle iterations...")
    start = (100, 100, 100)
    bi_diagram(beetle_func, arange(0, 1, 0.001), itemgetter(0), start)

def main():
    funcmap = dict(logistic=logistic_diagram,
                   beetle=beetle_diagram,
                   )
    def usage():
        print 'Usage: python %s [%s]' % (sys.argv[0], ' | '.join(funcmap.keys()))
    funckey = len(sys.argv) > 1 and sys.argv[1] or None
    func = funckey in funcmap and funcmap[funckey] or usage
    func()

if __name__ == '__main__':
    main()
